import cv2

def load_mnist(path, kind='train'):
    import os
    import gzip
    import numpy as np

    """Load MNIST data from `path`"""
    labels_path = os.path.join(path,
                               '%s-labels-idx1-ubyte'
                               % kind)
    images_path = os.path.join(path,
                               '%s-images-idx3-ubyte'
                               % kind)

    with open(labels_path, 'rb') as lbpath:
        labels = np.frombuffer(lbpath.read(), dtype=np.uint8,
                               offset=8)

    with open(images_path, 'rb') as imgpath:
        images = np.frombuffer(imgpath.read(), dtype=np.uint8,
                               offset=16).reshape(len(labels), 28, 28)

    return images, labels


def FMNIST(path):
    x_train, y_train = load_mnist(path, kind='train')
    x_test, y_test = load_mnist(path, kind='t10k')
    return (x_train,y_train),(x_test,y_test) 

        
def test():
    (x_train,y_train),(x_test,y_test) = FMNIST('../data/')
    print(x_train.shape,y_train.shape)
    print(x_test.shape,y_test.shape)
    print(y_train)
    print(y_test)
    cv2.imwrite('../data/x_train.jpg',x_train[0])
    cv2.imwrite('../data/x_test.jpg',x_test[0])


if __name__ == '__main__':
    test()